package cn.zhz.cloud.gateway.config;

import cn.zhz.cloud.common.model.result.R;
import cn.zhz.cloud.common.model.result.ResultCode;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.GatewayCallbackManager;
import com.alibaba.csp.sentinel.slots.system.SystemRule;
import com.alibaba.csp.sentinel.slots.system.SystemRuleManager;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.server.ServerResponse;

import java.util.ArrayList;
import java.util.List;

import static org.springframework.web.reactive.function.BodyInserters.fromValue;

/**
 * @author ZHZ
 * @since 2023-02-06
 */
@Configuration
public class SentinelConfig implements CommandLineRunner {

    @Override
    public void run(String... args) {
        List<SystemRule> ruleList = new ArrayList<>();
        SystemRule rule = new SystemRule();
        // 指定QPS限流阈值
        rule.setQps(1000);
        ruleList.add(rule);
        // 加载该规则
        SystemRuleManager.loadRules(ruleList);
        // 自定义限流返回信息
        GatewayCallbackManager.setBlockHandler((serverWebExchange, throwable) ->
                ServerResponse.status(HttpStatus.TOO_MANY_REQUESTS)
                        .contentType(MediaType.APPLICATION_JSON)
                        .body(fromValue(R.custom(ResultCode.TOO_MANY_REQUESTS, null))));
    }

}
